package com.wgh.demo.ai.controller;

import com.wgh.demo.ai.factory.AiModelFactory;
import com.wgh.demo.ai.service.AiModelService;
import com.wgh.demo.ai.service.ModelDecisionService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping("/api/ai/chat")
public class AiChatController {
    private final AiModelFactory aiModelFactory;
    private final ModelDecisionService modelDecisionService;

    @Autowired
    public AiChatController(AiModelFactory aiModelFactory, ModelDecisionService modelDecisionService) {
        this.aiModelFactory = aiModelFactory;
        this.modelDecisionService = modelDecisionService;
    }

    /**
     * 处理AI聊天请求
     * @param requestBody 请求体，包含用户输入的消息和可选的模型名称
     * @return 包含AI回复的响应
     */
    @PostMapping
    public Mono<ResponseEntity<Map<String, String>>> chat(@RequestBody Map<String, String> requestBody) {
        String message = requestBody.get("message");
        String modelName = requestBody.get("modelName");

        if (message == null || message.trim().isEmpty()) {
            return Mono.just(ResponseEntity.badRequest().body(Map.of("error", "消息不能为空")));
        }

        // 如果指定了模型名称，使用指定的模型
        if (modelName != null && !modelName.trim().isEmpty()) {
            AiModelService aiModelService = aiModelFactory.getAiModelService(modelName);
            // 调用AI模型生成回复
            return aiModelService.generateText(message)
                    .map(response -> ResponseEntity.ok(Map.of(
                            "response", response,
                            "modelUsed", aiModelService.getModelName()
                    )));
        } else {
            // 否则使用模型决策服务选择最合适的模型
            return modelDecisionService.decideModel(message)
                    .flatMap(decisionModelName -> {
                        AiModelService aiModelService = aiModelFactory.getAiModelService(decisionModelName);
                        // 调用AI模型生成回复
                        return aiModelService.generateText(message)
                                .map(response -> ResponseEntity.ok(Map.of(
                                        "response", response,
                                        "modelUsed", aiModelService.getModelName()
                                )));
                    });
        }
    }
}